package io.github.wanggit.antrpc.zipkin.dependencies.service.impl;

import com.alibaba.fastjson.JSONObject;
import io.github.wanggit.antrpc.zipkin.dependencies.constants.Constants;
import io.github.wanggit.antrpc.zipkin.dependencies.domain.Links;
import io.github.wanggit.antrpc.zipkin.dependencies.service.DependenciesService;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.aggregations.Aggregation;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.bucket.MultiBucketsAggregation;
import org.elasticsearch.search.aggregations.bucket.terms.ParsedStringTerms;
import org.elasticsearch.search.aggregations.metrics.ParsedValueCount;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.HashSet;
import java.util.Set;
import java.util.stream.Collectors;

@Slf4j
@Service
public class DependenciesServiceImpl implements DependenciesService {

    @Autowired private RestHighLevelClient restHighLevelClient;

    @Override
    public void apply(String yyyyMMdd) {
        if (null == yyyyMMdd) {
            yyyyMMdd = new SimpleDateFormat("yyyy-MM-dd").format(new Date());
        }
        if (log.isInfoEnabled()) {
            log.info("Calc Zipkin Dependencies.... yyyyMMdd=" + yyyyMMdd);
        }
        Set<String> allServiceNames = getAllServiceNames(yyyyMMdd);
        Set<Links> allLinks = getAllLinks(yyyyMMdd, allServiceNames);
        saveLinksToES(yyyyMMdd, allLinks);
    }

    private void saveLinksToES(String yyyyMMdd, Set<Links> allLinks) {
        for (Links link : allLinks) {
            String esDependencyIndex = Constants.ES_DEPENDENCY_INDEX_PREFIX + yyyyMMdd;
            IndexRequest indexRequest = new IndexRequest(esDependencyIndex);
            indexRequest.source(JSONObject.toJSONString(link), XContentType.JSON);
            try {
                IndexResponse indexResponse =
                        restHighLevelClient.index(indexRequest, RequestOptions.DEFAULT);
                int status = indexResponse.status().getStatus();
                System.out.println(status);
            } catch (IOException e) {
                if (log.isErrorEnabled()) {
                    log.error("Insert Links ERROR.", e);
                }
            }
        }
    }

    private Set<Links> getAllLinks(String yyyyMMdd, Set<String> allServiceNames) {
        Set<Links> linksSet = new HashSet<>();
        for (String parent : allServiceNames) {
            for (String child : allServiceNames) {
                SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
                searchSourceBuilder
                        .query(
                                QueryBuilders.boolQuery()
                                        .must(
                                                QueryBuilders.termQuery(
                                                        "localEndpoint.serviceName", parent))
                                        .must(
                                                QueryBuilders.termQuery(
                                                        "remoteEndpoint.serviceName", child)))
                        .fetchSource(false);
                searchSourceBuilder.aggregation(AggregationBuilders.count("id.count").field("id"));
                SearchRequest searchRequest =
                        new SearchRequest(Constants.ES_SPAN_INDEX_PREFIX + yyyyMMdd);
                searchRequest.source(searchSourceBuilder);
                try {
                    SearchResponse searchResponse =
                            restHighLevelClient.search(searchRequest, RequestOptions.DEFAULT);
                    Aggregation aggregation = searchResponse.getAggregations().get("id.count");
                    long count = ((ParsedValueCount) aggregation).getValue();
                    if (count > 0) {
                        Links links = new Links();
                        links.setCallCount(count);
                        links.setChild(child);
                        links.setErrorCount(0);
                        links.setParent(parent);
                        links.setId(parent + "|" + child);
                        linksSet.add(links);
                    }
                } catch (IOException e) {
                    if (log.isErrorEnabled()) {
                        log.error("Query Links ERROR.", e);
                    }
                }
            }
        }
        return linksSet;
    }

    private Set<String> getAllServiceNames(String yyyyMMdd) {
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.fetchSource(false);
        searchSourceBuilder
                .aggregation(
                        AggregationBuilders.terms("localEndpoint")
                                .field("localEndpoint.serviceName")
                                .size(1000))
                .aggregation(
                        AggregationBuilders.terms("remoteEndpoint")
                                .field("remoteEndpoint.serviceName")
                                .size(1000));

        SearchRequest searchRequest = new SearchRequest(Constants.ES_SPAN_INDEX_PREFIX + yyyyMMdd);
        searchRequest.source(searchSourceBuilder);
        Set<String> serviceNames = new HashSet<>();
        try {
            SearchResponse searchResponse =
                    restHighLevelClient.search(searchRequest, RequestOptions.DEFAULT);
            searchResponse
                    .getAggregations()
                    .asList()
                    .forEach(
                            it -> {
                                serviceNames.addAll(
                                        ((ParsedStringTerms) it)
                                                .getBuckets().stream()
                                                        .map(
                                                                MultiBucketsAggregation.Bucket
                                                                        ::getKeyAsString)
                                                        .collect(Collectors.toList()));
                            });
        } catch (IOException e) {
            if (log.isErrorEnabled()) {
                log.error("Query Service Names ERROR.", e);
            }
        }
        return serviceNames;
    }
}
